import os

import cv2
import numpy as np
from progressbar import *

from .module import DATASET
from .base_dataset import BaseDataset


@DATASET.register_module
class CUB(BaseDataset):

    def __init__(self, path, mode='train'):
        super(CUB, self).__init__(path, mode)

    def load_data(self):
        train_test_split_path = '{}/train_test_split.txt'.format(self.data_path)

        data_type = '1' if self.mode == 'train' else '0'
        with open(train_test_split_path, 'r') as f:
            line = f.readline()
            while line != '':
                line = line.strip('\n')
                id, mode = line.split(' ')
                line = f.readline()
                if mode == data_type:
                    self.cur_indexes.append(id)

        widgets = ['{} loading: '.format(self.__class__.__name__), Percentage(), ' ', Bar('|'),
                   ' ', Timer(), ' ', ETA(), ' ']
        count = 0
        pbar = ProgressBar(widgets=widgets, maxval=6000).start()

        images_path = '{}/images.txt'.format(self.data_path)
        img_dict = {}
        with open(images_path, 'r') as f:
            line = f.readline()
            while line != '':
                line = line.strip('\n')
                id, img_path = line.split(' ')
                img_dict[id] = img_path
                line = f.readline()
            f.close()
        for id in self.cur_indexes:
            self.img_paths.append(os.path.join('images', img_dict[id]))
            img = cv2.imread(os.path.join(self.data_path, 'images', img_dict[id]))
            self.sizes.append(np.array(img.shape[:-1]))
            count += 1
            pbar.update(count)
        pbar.finish()

        class_label_path = '{}/image_class_labels.txt'.format(self.data_path)
        cls_dict = {}
        with open(class_label_path, 'r') as f:
            line = f.readline()
            while line != '':
                line = line.strip('\n')
                id, cls = line.split(' ')
                line = f.readline()
                cls_dict[id] = int(cls)
        for id in self.cur_indexes:
            self.annotation.append(cls_dict[id] - 1)

        bounding_boxes_path = '{}/bounding_boxes.txt'.format(self.data_path)
        bounding_boxes_dict = {}
        with open(bounding_boxes_path, 'r') as f:
            line = f.readline()
            while line != '':
                line = line.strip('\n')
                id, x, y, w, h = line.split(' ')
                x = int(float(x))
                y = int(float(y))
                w = int(float(w))
                h = int(float(h))
                line = f.readline()
                bounding_boxes_dict[id] = [x, y, x+w, y+h]
        for i, id in enumerate(self.cur_indexes):
            # bounding_boxes_dict[id][2] = min(bounding_boxes_dict[id][2], self.sizes[i][1] - 1)
            # bounding_boxes_dict[id][3] = min(bounding_boxes_dict[id][3], self.sizes[i][0] - 1)
            self.annotation[i] = np.array([self.annotation[i]] + bounding_boxes_dict[id], dtype=np.int32)

        print('{}: Loaded {:d} samples for {}'
              .format(self.__class__.__name__, len(self.cur_indexes), self.mode))